{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "华为美食比赛的代码"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 一、库导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load AI/code/lib/imports.py\n",
    "# %load AI/code/lib/imports.py\n",
    "\n",
    "# 系统库\n",
    "import os\n",
    "import time\n",
    "import pathlib\n",
    "import argparse\n",
    "\n",
    "# 数据图像处理\n",
    "import numpy as np \n",
    "import pandas as pd \n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# PyTorch框架\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import transforms, models\n",
    "import torchvision.datasets as datasets\n",
    "from torch.utils.data import SubsetRandomSampler\n",
    "from torch.autograd import Variable\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import kaggle  # AI比赛数据下载\n",
    "import ttach as tta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 二、环境配置"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## jupyter使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# %%writefile AI/code/help/jupyter.md\n",
    "# %load AI/code/help/jupyter.md\n",
    "# jupyter notebook使用\n",
    "\n",
    "# # 查看变量\n",
    "# ?traindata\n",
    "\n",
    "# # 查看库文件内容\n",
    "# ??os\n",
    "\n",
    "# # 查看类\n",
    "# ???dta\n",
    "\n",
    "# # 删除一个变量\n",
    "# %xdel model\n",
    "\n",
    "# # 自动导入\n",
    "# %reload_ext autoreload\n",
    "# %autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple\n",
      "Collecting ttach\n",
      "  Downloading https://pypi.tuna.tsinghua.edu.cn/packages/53/22/470bb42f90505dc572f6bbcf3ac84d67aaf1554cd48cc08f788c36fec129/ttach-0.0.2-py3-none-any.whl (9.1 kB)\n",
      "Installing collected packages: ttach\n",
      "Successfully installed ttach-0.0.2\n"
     ]
    }
   ],
   "source": [
    "!pip install ttach"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 自动导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 参数配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load AI/code/config/parameter.py\n",
    "parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\n",
    "parser.add_argument('--data', metavar='DIR',\n",
    "                    help='数据的路径')\n",
    "parser.add_argument('--j', '--workers', default=4, type=int, metavar='N',\n",
    "                    help='数据导入的线程 (默认为: 4)')\n",
    "parser.add_argument('--epochs', default=90, type=int, metavar='N',\n",
    "                    help='训练的次数')\n",
    "parser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n",
    "                    help='开始的次数 (useful on restarts)')\n",
    "parser.add_argument('--b', '--batch-size', default=16, type=int,\n",
    "                    metavar='N',\n",
    "                    help='批处理大小（默认为16）')\n",
    "parser.add_argument('--lr', '--learning-rate', default=[5e-4, 1e-4, 1e-5, 1e-6], type=float,\n",
    "                    metavar='LR', help='初始化的学习率', dest='lr')\n",
    "parser.add_argument('--setlr', '--learning-change', default=[0, 5, 9], type=int,\n",
    "                     help='初始化的学习率', dest='setlr')\n",
    "parser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n",
    "                    help='动量')\n",
    "parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,\n",
    "                    metavar='W', help='权重下降 (默认: 1e-4)',\n",
    "                    dest='weight_decay')\n",
    "parser.add_argument('--p', '--print-freq', default=10, type=int,\n",
    "                    metavar='N', help='打印频率 (默认: 10)')\n",
    "parser.add_argument('--resume', default=None, type=str, metavar='PATH',\n",
    "                    help='恢复的模型 (默认: 空)')\n",
    "parser.add_argument('--e', '--evaluate', dest='evaluate', action='store_true',\n",
    "                    help='evaluate model on validation set')\n",
    "parser.add_argument('--pretrained', dest='pretrained', action='store_true',default=True,\n",
    "                    help='使用训练模型')\n",
    "parser.add_argument('--world-size', default=-1, type=int,\n",
    "                    help='分布式训练的节点数')\n",
    "parser.add_argument('--rank', default=-1, type=int,\n",
    "                    help='分布式训练的节点等级')\n",
    "parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,\n",
    "                    help='用于设置分布式培训的url')\n",
    "parser.add_argument('--dist-backend', default='nccl', type=str,\n",
    "                    help='分布式后端')\n",
    "parser.add_argument('--seed', default=None, type=int,\n",
    "                    help='初始化培训的种子. ')\n",
    "parser.add_argument('--gpu', default=0, type=int,\n",
    "                    help='GPU使用的ID.')\n",
    "parser.add_argument('--multiprocessing-distributed', action='store_true',\n",
    "                    help='使用多处理分布式培训来启动每个节点有N个进程，'\n",
    "                         '其中有N个GPU。这是对单个节点或PyTorch使用'\n",
    "                         'PyTorch的最快方法多节点数据并行训练')\n",
    "\n",
    "parser.add_argument('--num_classes', type=int, default=10, help='您的任务应该分类的类数')\n",
    "parser.add_argument('--local_data_root', default='/cache/', type=str,\n",
    "                    help='a directory used for transfer data between local path and OBS path')\n",
    "parser.add_argument('--data_url', type=str, help='the training and validation data path')\n",
    "parser.add_argument('--test_data_url', default='', type=str, help='the test data path')\n",
    "parser.add_argument('--data_local', default='', type=str, help='the training and validation data path on local')\n",
    "parser.add_argument('--test_data_local', default='', type=str, help='the test data path on local')\n",
    "parser.add_argument('--train_url', type=str, help='the path to save training outputs')\n",
    "parser.add_argument('--train_local', default='', type=str, help='the training output results on local')\n",
    "parser.add_argument('--tmp', default='', type=str, help='a temporary path on local')\n",
    "parser.add_argument('--deploy_script_path', default='', type=str,\n",
    "                    help='a path which contain config.json and customize_service.py, '\n",
    "                         'if it is set, these two scripts will be copied to {train_url}/model directory')\n",
    "\n",
    "args = parser.parse_args(args=[]) # 实际使用要删除\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args.pretrained"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%writefile AI/code/dataprocess/Food_Category.py\n",
    "from AI.code.utils.util import photo_process\n",
    "import os,shutil\n",
    "food_path = \"I:\\\\AI\\\\数据集\\\\美食分类\\\\data\\\\images\"\n",
    "train_data_path = \"data\\\\goodeat\\\\train\"\n",
    "val_data_path   = \"data\\\\goodeat\\\\val\"\n",
    "test_data_path = \"data\\\\goodeat\\\\test\"\n",
    "\n",
    "traintest_cent = 0.8\n",
    "trainval_cent = 0.75\n",
    "\n",
    "# 图片处理\n",
    "# photo_process(food_path, train_data_path, val_data_path, test_data_path, traintest_cent, trainval_cent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据加载器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%writefile AI/code/dataloader/path_dataloader.py\n",
    "# %load AI/code/dataloader/data_dataloader.py\n",
    "# 制作数据的数据导入器\n",
    "\n",
    "train_data_transform = transforms.Compose([\n",
    "#     transforms.Resize(256),\n",
    "#     transforms.CenterCrop(224),\n",
    "    transforms.Resize((224,224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "val_data_transform = transforms.Compose([\n",
    "#     transforms.Resize(256),\n",
    "#     transforms.CenterCrop(224),\n",
    "    transforms.Resize((224,224)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "test_data_transform = transforms.Compose([\n",
    "    transforms.Resize((288,288)),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "# 设置batch大小\n",
    "train_batch_size = 8\n",
    "val_batch_size = 8\n",
    "test_batch_size = 1\n",
    "\n",
    "# DataSet类 - （目录中图片导入）\n",
    "train_dataset = datasets.ImageFolder(root=train_data_path, transform=train_data_transform)\n",
    "val_dataset = datasets.ImageFolder(root=val_data_path, transform=val_data_transform)\n",
    "test_dataset = datasets.ImageFolder(root=val_data_path, transform=val_data_transform)\n",
    "\n",
    "# 数据加载器\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=0)\n",
    "val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=val_batch_size, shuffle=True, num_workers=0)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(val_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'三明治': 0,\n",
       " '冰激凌': 1,\n",
       " '土豆泥': 2,\n",
       " '小米粥': 3,\n",
       " '松鼠鱼': 4,\n",
       " '烤冷面': 5,\n",
       " '玉米饼': 6,\n",
       " '甜甜圈': 7,\n",
       " '芒果班戟': 8,\n",
       " '鸡蛋布丁': 9}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset.class_to_idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型建立"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load AI/code/models/classifier.py\n",
    "import os,sys\n",
    "from torchvision import models\n",
    "from AI.models.efficientnet import EfficientNet\n",
    "from AI.models.resnetxt_wsl import resnext101_32x8d_wsl,resnext101_32x16d_wsl,resnext101_32x32d_wsl,resnext101_32x48d_wsl\n",
    "import pretrainedmodels\n",
    "\n",
    "# 以脚本执行时允许相对导入\n",
    "if __name__ == \"__main__\" and __package__ is None:\n",
    "    sys.path.insert(0, os.path.join(os.getcwd(), '..', '..', '..'))\n",
    "    import AI  # noqa: F401\n",
    "    __package__ = \"AI\"\n",
    "\n",
    "from AI.models.classifier import  Classifier\n",
    "\n",
    "def set_model():\n",
    "    if args.pretrained:\n",
    "        print(\"Model have pretrained model\")\n",
    "#         model = models.resnet18(pretrained=True)\n",
    "    #     model = models.resnet50(pretrained=True)\n",
    "    #     model = models.resnet101(pretrained=True)\n",
    "    #     model = models.resnet152(pretrained=True)\n",
    "    #     model = models.resnext50_32x4d(pretrained=True)\n",
    "        model = models.resnext101_32x4d(pretrained=True)\n",
    "    #     model = resnext101_32x8d_wsl()\n",
    "    #     model = resnext101_32x16d_wsl()\n",
    "    #     model = resnext101_32x32d_wsl()\n",
    "    #     model = resnext101_32x48d_wsl()\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b0') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b1') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b2') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b3') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b4') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b5') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b6') # 加载预训练模型\n",
    "    #     model = EfficientNet.from_pretrained('efficientnet-b7') # 加载预训练模型\n",
    "        num_ftrs = model.fc.in_features\n",
    "        model.fc = nn.Sequential(\n",
    "            nn.Linear(num_ftrs, 1024),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1024, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 10),\n",
    "            nn.LogSoftmax(dim=1),\n",
    "        )  \n",
    "\n",
    "    else:\n",
    "        print(\"Model isn't have pretrained model\")\n",
    "    #     model = EfficientNet.from_name('efficientnet-b0') # 不加载预训练模型\n",
    "    #     model = Classifier()\n",
    "        # model = models.resnet50(pretrained=True)\n",
    "    #     model = models.resnet18(pretrained=True)\n",
    "        model = models.resnext50_32x4d(pretrained=True)\n",
    "        num_ftrs = model._fc.in_features\n",
    "        model._fc = nn.Sequential(\n",
    "            nn.Linear(num_ftrs, 1024),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1024, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 10),\n",
    "            nn.LogSoftmax(dim=1),\n",
    "        )\n",
    "\n",
    "    if args.gpu > -1:\n",
    "        model = model.cuda()\n",
    "    else:\n",
    "        model = model.cpu()\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pretrainedmodels.model_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %xdel Classifier\n",
    "# %xdel model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 10])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(torch.rand(1,3,224,224).cuda()).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load AI/code/models/criterions.py\n",
    "# %load AI/code/models/criterions.py\n",
    "# 损失函数\n",
    "# 多分类（cross entropy）\n",
    "\n",
    "# 多标签分类（binary cross entropy）\n",
    "# 语义分割（binary cross entropy\n",
    "# 如果是多类语义分割，可选cross entropy）\n",
    "\n",
    "# criterion = nn.NLLLoss()\n",
    "def set_criterion():\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    return criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load AI/code/models/optimizers.py\n",
    "# 优化器\n",
    "\n",
    "# 常用Adam优化器， 细致调参时再用SGD\n",
    "# Adam: init_lr=5e-4(3e-4)    \n",
    "# 衰减步骤 [5e-4(3e-4), 1e-4, 1e-5, 1e-6]\n",
    "# optimizer = torch.optim.SGD(model.parameters(), args.lr[0],\n",
    "#                                 momentum=args.momentum,\n",
    "#                                 weight_decay=args.weight_decay)\n",
    "# optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)\n",
    "# optimizer = optim.Adam([var1, var2], lr = 0.0001)\n",
    "# optim.SGD([\n",
    "#                 {'params': model.base.parameters()},\n",
    "#                 {'params': model.classifier.parameters(), 'lr': 1e-3}\n",
    "#             ], lr=1e-2, momentum=0.9)\n",
    "# optimizer = torch.optim.Adam([{'params': model.backbone.parameters(), 'lr': 3e-5},\n",
    "#                               {'params': model.fc.parameters(), 'lr': 3e-4}, ])\n",
    "# optimizer = optim.Adam(model.parameters(), lr=args.lr[0])\n",
    "# params = list(model.fc.parameters())\n",
    "# optimizer = optim.SGD(params, lr=0.001, momentum=0.9)\n",
    "def set_optimizer(model):\n",
    "#     optimizer = optim.SGD([\n",
    "#                     {'params': model.layer1.parameters()},\n",
    "#                     {'params': model.layer2.parameters()},\n",
    "#                     {'params': model.layer3.parameters()},\n",
    "#                     {'params': model.layer4.parameters()},\n",
    "#                     {'params': model.avgpool.parameters()},\n",
    "#                 ], lr=0.0001, momentum=0.9)\n",
    "    optimizer = optim.SGD(model.parameters(), lr = 0.00005, momentum=0.9)\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型训练与验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# %load AI/code/trainval/class_train.py\n",
    "# %load AI/code/trainval/class_train.py\n",
    "# %load AI/code/trainval/class_train.py\n",
    "# %load AI/code/trainval/class_train.py\n",
    "from AI.utils.train_utils import get_accuracy, save_checkpoint, logger, adjust_learning_rate, AverageMeter\n",
    "\n",
    "def main(run_num):\n",
    "    train_logger, val_logger = main_worker(run_num, gpu_flag=True)\n",
    "    return train_logger, val_logger\n",
    "\n",
    "def main_worker(run_num, gpu_flag=False):\n",
    "    train_log = logger('runs//'+str(run_num)+ '//train')\n",
    "    val_log = logger('runs//'+str(run_num)+'//val')\n",
    "    \n",
    "    # 设置损失函数\n",
    "    criterion = set_criterion()\n",
    "    optimizer = set_optimizer(model)\n",
    "\n",
    "    if args.resume:\n",
    "        print(\"=> loading checkpoint '{}'\".format(args.resume))\n",
    "        if args.gpu is None:\n",
    "            checkpoint = torch.load(args.resume)\n",
    "        else:\n",
    "            # Map model to be loaded to specified single gpu.\n",
    "            loc = 'cuda:{}'.format(args.gpu)\n",
    "            checkpoint = torch.load(args.resume, map_location=loc)\n",
    "        args.start_epoch = checkpoint['epoch']\n",
    "        best_acc1 = checkpoint['best_acc1']    \n",
    "        model.load_state_dict(checkpoint['state_dict'])\n",
    "        # optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "        print(\"=> loaded checkpoint '{}' (epoch {})\"\n",
    "              .format(args.resume, checkpoint['epoch']))\n",
    "        \n",
    "    train_logger, val_logger = [[],[],[]], [[],[],[]]\n",
    "    best_acc1 = 0\n",
    "    tmp_time=start_time = time.time()\n",
    "    # 开始训练\n",
    "    print('Start training!')\n",
    "    for epoch in range(args.start_epoch, args.epochs):\n",
    "        adjust_learning_rate(optimizer, epoch, args.lr, args.setlr)\n",
    "        \n",
    "        train_out = train_model(epoch, train_loader, model, criterion, optimizer, train_log, gpu_flag)\n",
    "        val_out = val_model(epoch, val_loader, model, criterion, optimizer, val_log, gpu_flag)\n",
    "        train_log.add_train_data(train_out, epoch)\n",
    "        val_log.add_train_data(val_out, epoch)\n",
    "        is_best = val_out[1] > best_acc1\n",
    "        best_acc1 = max(val_out[1], best_acc1)\n",
    "        save_checkpoint({\n",
    "                'epoch': epoch + 1,\n",
    "                'arch': model,\n",
    "                'state_dict': model.state_dict(),\n",
    "                'best_acc1' : best_acc1,\n",
    "                'acc1': val_out[1],\n",
    "                'acc5': val_out[2],\n",
    "                'optimizer' : optimizer.state_dict(),\n",
    "            },is_best, 'resnet50')\n",
    "\n",
    "        for i in range(3):\n",
    "            train_logger[i].append(train_out[i])\n",
    "            val_logger[i].append(val_out[i])\n",
    "\n",
    "        print(\"Epoch: {epoch} Tloss：{loss:6.4f}    Vloss:{vloss:6.4f} \\t Acc@1_avg {acc1:6.2f} \\t Acc@5_avg {acc5:6.2f} \\t \\\n",
    "Time@ {time:6.2f}\\t TotalTime@{ttime:6.2f}\"\n",
    "                                                          .format(epoch=epoch,  \n",
    "                                                                  loss=train_out[0],         vloss=val_out[0],\n",
    "                                                                  acc1=val_out[1],           acc5=val_out[2],\n",
    "                                                                  time=time.time()-tmp_time, ttime=time.time()-start_time))\n",
    "        tmp_time = time.time()\n",
    "    print('Finished training!')\n",
    "    train_log.close(),val_log.close()\n",
    "    \n",
    "    return train_logger, val_logger\n",
    "    \n",
    "def train_model(epoch, data_loader, model, criterion, optimizer, log, gpu_flag=False):\n",
    "    '''\n",
    "    模型的训练函数\n",
    "    '''\n",
    "    if gpu_flag:\n",
    "        model.cuda().train()\n",
    "    else:\n",
    "        model.cpu().train()\n",
    "        \n",
    "    loss_avg = AverageMeter()\n",
    "    acc1_avg= AverageMeter()\n",
    "    acc5_avg= AverageMeter()\n",
    "    data_len = len(data_loader)\n",
    "    for i, data in enumerate(data_loader):\n",
    "        images, labels = data\n",
    "        if gpu_flag:\n",
    "            images = Variable(images).cuda()\n",
    "            labels = Variable(labels).cuda()\n",
    "        else:\n",
    "            images = Variable(images)\n",
    "            labels = Variable(labels)\n",
    "        # if i%2 == 1:\n",
    "        optimizer.zero_grad()\n",
    "        log_ps = model(images)\n",
    "        loss = criterion(log_ps, labels)\n",
    "        loss.backward()\n",
    "        # if i%2 == 1:\n",
    "        optimizer.step()\n",
    "        acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))\n",
    "        acc1_avg.update(acc1.item(), images.size(0))\n",
    "        acc5_avg.update(acc5.item(), images.size(0))\n",
    "        loss_avg.update(loss.item(), images.size(0))\n",
    "        log.add_train_in_data(loss.item(), data_len*epoch + i)\n",
    "        \n",
    "        if i%5000 == 4999:\n",
    "            print(\"Train Epoch: {epoch} Loss_avg：{loss:6.4f} \\t Acc@1_avg {acc1:6.2f} \\t Acc@5_avg {acc5:6.2f}\".format(epoch=epoch,\n",
    "                                                                  loss=loss_avg.avg,\n",
    "                                                                  acc1=acc1_avg.avg,\n",
    "                                                                  acc5=acc5_avg.avg))\n",
    "    \n",
    "    return loss_avg.avg, acc1_avg.avg, acc5_avg.avg\n",
    "        \n",
    "    \n",
    "def val_model(epoch, data_loader, model, criterion, optimizer, log, gpu_flag=False):\n",
    "    '''\n",
    "    验证函数\n",
    "    '''\n",
    "    if gpu_flag:\n",
    "        model.cuda().eval()\n",
    "    else:\n",
    "        model.cpu().eval()\n",
    "    \n",
    "    loss_avg = AverageMeter()\n",
    "    acc1_avg= AverageMeter()\n",
    "    acc5_avg= AverageMeter()\n",
    "    for i, data in enumerate(data_loader):\n",
    "        images, labels = data\n",
    "        if gpu_flag:\n",
    "            images = Variable(images).cuda()\n",
    "            labels = Variable(labels).cuda()\n",
    "        else:\n",
    "            images = Variable(images)\n",
    "            labels = Variable(labels)\n",
    "        log_ps = model(images)\n",
    "        loss = criterion(log_ps, labels)\n",
    "        acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))\n",
    "        acc1_avg.update(acc1.item(), images.size(0))\n",
    "        acc5_avg.update(acc5.item(), images.size(0))\n",
    "        loss_avg.update(loss.item(), images.size(0))\n",
    "        \n",
    "        if i%5000 == 4999:\n",
    "            print(\"Val   Epoch: {epoch} Loss_avg：{loss:6.4f} \\t Acc@1_avg {acc1:6.2f} \\t Acc@5_avg {acc5:6.2f}\".format(epoch=epoch,\n",
    "                                                                  loss=loss_avg.avg,\n",
    "                                                                  acc1=acc1_avg.avg,\n",
    "                                                                  acc5=acc5_avg.avg))\n",
    "    \n",
    "    return loss_avg.avg, acc1_avg.avg, acc5_avg.avg\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start training!\n",
      "epoch  0 set lr is  0.001\n",
      "Epoch: 0 Tloss：0.1594    Vloss:0.2870 \t Acc@1_avg  92.00 \t Acc@5_avg  99.50 \t Time@  72.96\t TotalTime@ 72.96\n",
      "Finished training!\n"
     ]
    }
   ],
   "source": [
    "# 训练参数定义\n",
    "args.start_epoch = 0\n",
    "args.epochs = 1\n",
    "args.lr = [0.001, 0.001, 0.001, 0.001]\n",
    "args.setlr = [0,5,29]\n",
    "# args.resume=\"checkpoints//5_resnet50_95.54.pth\"\n",
    "# args.resume=\"checkpoints//model_best.pth\"\n",
    "# args.resume=None\n",
    "\n",
    "# 训练主函数                             运行次数\n",
    "train_logger, val_logger = main(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# %%writefile AI/code/utils/loss_show.py\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "plt.figure(figsize=(15,4))\n",
    "plt.title(\"342342\")\n",
    "plt.subplot(131) \n",
    "plt.plot(train_logger[0], label='Train loss')\n",
    "plt.plot(val_logger[0], label='Val loss')\n",
    "plt.legend(frameon=False)\n",
    "plt.subplot(132) \n",
    "plt.plot(train_logger[1], label='Train acc1')\n",
    "plt.plot(val_logger[1], label='Val acc1')\n",
    "plt.legend(frameon=False)\n",
    "plt.title('Radio',fontsize=18,color='r')\n",
    "plt.subplot(133)\n",
    "plt.plot(train_logger[2], label='Train acc5')\n",
    "plt.plot(val_logger[2], label='Val acc5')\n",
    "plt.legend(frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def testmodel(data_loader, model, num, gpu_flag=False):\n",
    "    if gpu_flag:\n",
    "        model.cuda().train()\n",
    "    else:\n",
    "        model.cpu().train()\n",
    "        \n",
    "    for i, data in enumerate(data_loader):\n",
    "        if i == num:\n",
    "            images, labels = data\n",
    "            if gpu_flag:\n",
    "                images = Variable(images).cuda()\n",
    "                labels = Variable(labels).cuda()\n",
    "            else:\n",
    "                images = Variable(images)\n",
    "                labels = Variable(labels)\n",
    "            log_ps = model(images)\n",
    "            ps = torch.exp(log_ps)\n",
    "            confid, classes = ps.topk(1)\n",
    "            if classes.item() == labels.item():\n",
    "                print(\"success\")\n",
    "            else:\n",
    "                print( classes.item(), labels.item(), confid.item())\n",
    "                print(torch.max(images), torch.min(torch.min(images)))\n",
    "                images = images*0.5 + 1\n",
    "                img = images.squeeze(0).cpu().numpy().transpose(1,2,0)\n",
    "                plt.imshow(img)\n",
    "                plt.show()\n",
    "        \n",
    "        \n",
    "    return \n",
    "\n",
    "# testmodel(test_loader,model, 3,True)\n",
    "test_num = len(test_loader)\n",
    "num = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testmodel(test_loader,model, num,True)\n",
    "num = num + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型调优"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据增强"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tun_data_path = \"data\\\\goodeat\\\\train\"\n",
    "\n",
    "transform_c = [\n",
    "    transforms.RandomHorizontalFlip(p=0.5),  # 水平翻转\n",
    "    transforms.RandomVerticalFlip(p=0.5),    # 竖直翻转\n",
    "    transforms.ColorJitter(brightness=(0, 36), contrast=(\n",
    "        0, 10), saturation=(0, 25), hue=(-0.1, 0.1)),\n",
    "    # transforms.RandomPerspective(distortion_scale=1, p=1, interpolation=3),  # 透视变换\n",
    "    transforms.RandomAffine(degrees=30, translate=(0, 0.2), scale=(0.9, 1), shear=(6, 9), fillcolor=66) # 中间旋转\n",
    "]\n",
    "\n",
    "tun_data_transform = transforms.Compose([\n",
    "    transforms.Resize((224,224)),\n",
    "#     transforms.RandomChoice(transform_c), # 随机从一组变换中选择一个\n",
    "#     transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "tun_dataset = datasets.ImageFolder(root=tun_data_path, transform=tun_data_transform)\n",
    "tun_loader = torch.utils.data.DataLoader(tun_dataset, batch_size=1, shuffle=True, num_workers=0)\n",
    "\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.ToPILImage() # Convert a tensor or an ndarray to PIL Image.\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = 'H:\\\\Pytorch-RetinaNet\\\\data\\\\goodeat\\\\train\\\\芒果班戟'\n",
    "filelist = os.listdir(path)\n",
    "img = Image.open(os.path.join(path, filelist[0]))\n",
    "img = tun_data_transform(img)\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for image, label in tun_loader:\n",
    "    img = image.squeeze(0).numpy().transpose(1,2,0)\n",
    "    plt.imshow(img)\n",
    "    plt.show()\n",
    "    print(img.shape)\n",
    "    \n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "IndentationError",
     "evalue": "expected an indented block (<ipython-input-20-da1414f2ccc6>, line 3)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  File \u001b[1;32m\"<ipython-input-20-da1414f2ccc6>\"\u001b[1;36m, line \u001b[1;32m3\u001b[0m\n\u001b[1;33m    criterion = set_criterion()\u001b[0m\n\u001b[1;37m            ^\u001b[0m\n\u001b[1;31mIndentationError\u001b[0m\u001b[1;31m:\u001b[0m expected an indented block\n"
     ]
    }
   ],
   "source": [
    "from AI.utils.train_utils import get_accuracy, save_checkpoint, logger, adjust_learning_rate\n",
    "def test_model(model):\n",
    "    criterion = set_criterion()\n",
    "    optimizer = set_optimizer(model)\n",
    "    for i, data in enumerate(val_loader):\n",
    "            images, labels = data\n",
    "            if True:\n",
    "                images = Variable(images).cuda()\n",
    "                labels = Variable(labels).cuda()\n",
    "            else:\n",
    "                images = Variable(images)\n",
    "                labels = Variable(labels)\n",
    "            # if i%2 == 1:\n",
    "            optimizer.zero_grad()\n",
    "            log_ps = model(images)\n",
    "            loss = criterion(log_ps, labels)\n",
    "    #         loss.backward()\n",
    "            # if i%2 == 1:\n",
    "    #         optimizer.step()\n",
    "            print(log_ps.shape)\n",
    "            print(\"loss\", loss)\n",
    "            acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))\n",
    "            _, preds = torch.max(log_ps.data, 1)\n",
    "            tt = torch.sum(preds == labels.data).to(torch.float32)\n",
    "            print(tt)\n",
    "            print(acc1)\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TTA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "modelpath = \"H:\\迅雷下载\\model_101.pth\"\n",
    "model = models.resnext101_32x8d(pretrained=False)\n",
    "save_model = torch.load(modelpath)\n",
    "model_dict =  model.state_dict()\n",
    "state_dict = {k:v for k,v in save_model.items() if k in model_dict.keys()}\n",
    "model_dict.update(state_dict)\n",
    "model.load_state_dict(model_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "transforms = tta.Compose(\n",
    "    [\n",
    "        tta.HorizontalFlip(),\n",
    "        tta.Rotate90(angles=[0, 180]),\n",
    "        tta.Scale(scales=[1, 2, 4]),\n",
    "        tta.Multiply(factors=[0.9, 1, 1.1]),        \n",
    "    ]\n",
    ")\n",
    "\n",
    "tta_model = tta.ClassificationTTAWrapper(model, transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on class ClassificationTTAWrapper in module ttach.wrappers:\n",
      "\n",
      "class ClassificationTTAWrapper(torch.nn.modules.module.Module)\n",
      " |  ClassificationTTAWrapper(model: torch.nn.modules.module.Module, transforms: ttach.base.Compose, merge_mode: str = 'mean', output_label_key: Union[str, NoneType] = None)\n",
      " |  \n",
      " |  Wrap PyTorch nn.Module (classification model) with test time augmentation transforms\n",
      " |  \n",
      " |  Args:\n",
      " |      model (torch.nn.Module): classification model with single input and single output\n",
      " |          (.forward(x) should return either torch.Tensor or Mapping[str, torch.Tensor])\n",
      " |      transforms (ttach.Compose): composition of test time transforms\n",
      " |      merge_mode (str): method to merge augmented predictions mean/gmean/max/min/sum/tsharpen\n",
      " |      output_mask_key (str): if model output is `dict`, specify which key belong to `label`\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      ClassificationTTAWrapper\n",
      " |      torch.nn.modules.module.Module\n",
      " |      builtins.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  __init__(self, model: torch.nn.modules.module.Module, transforms: ttach.base.Compose, merge_mode: str = 'mean', output_label_key: Union[str, NoneType] = None)\n",
      " |      Initializes internal Module state, shared by both nn.Module and ScriptModule.\n",
      " |  \n",
      " |  forward(self, image: torch.Tensor, *args) -> Union[torch.Tensor, Mapping[str, torch.Tensor]]\n",
      " |      Defines the computation performed at every call.\n",
      " |      \n",
      " |      Should be overridden by all subclasses.\n",
      " |      \n",
      " |      .. note::\n",
      " |          Although the recipe for forward pass needs to be defined within\n",
      " |          this function, one should call the :class:`Module` instance afterwards\n",
      " |          instead of this since the former takes care of running the\n",
      " |          registered hooks while the latter silently ignores them.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __call__(self, *input, **kwargs)\n",
      " |      Call self as a function.\n",
      " |  \n",
      " |  __delattr__(self, name)\n",
      " |      Implement delattr(self, name).\n",
      " |  \n",
      " |  __dir__(self)\n",
      " |      Default dir() implementation.\n",
      " |  \n",
      " |  __getattr__(self, name)\n",
      " |  \n",
      " |  __repr__(self)\n",
      " |      Return repr(self).\n",
      " |  \n",
      " |  __setattr__(self, name, value)\n",
      " |      Implement setattr(self, name, value).\n",
      " |  \n",
      " |  __setstate__(self, state)\n",
      " |  \n",
      " |  add_module(self, name, module)\n",
      " |      Adds a child module to the current module.\n",
      " |      \n",
      " |      The module can be accessed as an attribute using the given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the child module. The child module can be\n",
      " |              accessed from this module using the given name\n",
      " |          module (Module): child module to be added to the module.\n",
      " |  \n",
      " |  apply(self, fn)\n",
      " |      Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\n",
      " |      as well as self. Typical use includes initializing the parameters of a model\n",
      " |      (see also :ref:`nn-init-doc`).\n",
      " |      \n",
      " |      Args:\n",
      " |          fn (:class:`Module` -> None): function to be applied to each submodule\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> def init_weights(m):\n",
      " |          >>>     print(m)\n",
      " |          >>>     if type(m) == nn.Linear:\n",
      " |          >>>         m.weight.data.fill_(1.0)\n",
      " |          >>>         print(m.weight)\n",
      " |          >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))\n",
      " |          >>> net.apply(init_weights)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 1.,  1.],\n",
      " |                  [ 1.,  1.]])\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 1.,  1.],\n",
      " |                  [ 1.,  1.]])\n",
      " |          Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |          Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |  \n",
      " |  buffers(self, recurse=True)\n",
      " |      Returns an iterator over module buffers.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          torch.Tensor: module buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for buf in model.buffers():\n",
      " |          >>>     print(type(buf.data), buf.size())\n",
      " |          <class 'torch.FloatTensor'> (20L,)\n",
      " |          <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  children(self)\n",
      " |      Returns an iterator over immediate children modules.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a child module\n",
      " |  \n",
      " |  cpu(self)\n",
      " |      Moves all model parameters and buffers to the CPU.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  cuda(self, device=None)\n",
      " |      Moves all model parameters and buffers to the GPU.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing optimizer if the module will\n",
      " |      live on GPU while being optimized.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  double(self)\n",
      " |      Casts all floating point parameters and buffers to ``double`` datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  eval(self)\n",
      " |      Sets the module in evaluation mode.\n",
      " |      \n",
      " |      This has any effect only on certain modules. See documentations of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  extra_repr(self)\n",
      " |      Set the extra representation of the module\n",
      " |      \n",
      " |      To print customized extra information, you should reimplement\n",
      " |      this method in your own modules. Both single-line and multi-line\n",
      " |      strings are acceptable.\n",
      " |  \n",
      " |  float(self)\n",
      " |      Casts all floating point parameters and buffers to float datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  half(self)\n",
      " |      Casts all floating point parameters and buffers to ``half`` datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  load_state_dict(self, state_dict, strict=True)\n",
      " |      Copies parameters and buffers from :attr:`state_dict` into\n",
      " |      this module and its descendants. If :attr:`strict` is ``True``, then\n",
      " |      the keys of :attr:`state_dict` must exactly match the keys returned\n",
      " |      by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          state_dict (dict): a dict containing parameters and\n",
      " |              persistent buffers.\n",
      " |          strict (bool, optional): whether to strictly enforce that the keys\n",
      " |              in :attr:`state_dict` match the keys returned by this module's\n",
      " |              :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
      " |      \n",
      " |      Returns:\n",
      " |          ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
      " |              * **missing_keys** is a list of str containing the missing keys\n",
      " |              * **unexpected_keys** is a list of str containing the unexpected keys\n",
      " |  \n",
      " |  modules(self)\n",
      " |      Returns an iterator over all modules in the network.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a module in the network\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.modules()):\n",
      " |                  print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |          1 -> Linear(in_features=2, out_features=2, bias=True)\n",
      " |  \n",
      " |  named_buffers(self, prefix='', recurse=True)\n",
      " |      Returns an iterator over module buffers, yielding both the\n",
      " |      name of the buffer as well as the buffer itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all buffer names.\n",
      " |          recurse (bool): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, torch.Tensor): Tuple containing the name and buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, buf in self.named_buffers():\n",
      " |          >>>    if name in ['running_var']:\n",
      " |          >>>        print(buf.size())\n",
      " |  \n",
      " |  named_children(self)\n",
      " |      Returns an iterator over immediate children modules, yielding both\n",
      " |      the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Module): Tuple containing a name and child module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, module in model.named_children():\n",
      " |          >>>     if name in ['conv4', 'conv5']:\n",
      " |          >>>         print(module)\n",
      " |  \n",
      " |  named_modules(self, memo=None, prefix='')\n",
      " |      Returns an iterator over all modules in the network, yielding\n",
      " |      both the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Module): Tuple of name and module\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.named_modules()):\n",
      " |                  print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> ('', Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          ))\n",
      " |          1 -> ('0', Linear(in_features=2, out_features=2, bias=True))\n",
      " |  \n",
      " |  named_parameters(self, prefix='', recurse=True)\n",
      " |      Returns an iterator over module parameters, yielding both the\n",
      " |      name of the parameter as well as the parameter itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all parameter names.\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Parameter): Tuple containing the name and parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, param in self.named_parameters():\n",
      " |          >>>    if name in ['bias']:\n",
      " |          >>>        print(param.size())\n",
      " |  \n",
      " |  parameters(self, recurse=True)\n",
      " |      Returns an iterator over module parameters.\n",
      " |      \n",
      " |      This is typically passed to an optimizer.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Parameter: module parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for param in model.parameters():\n",
      " |          >>>     print(type(param.data), param.size())\n",
      " |          <class 'torch.FloatTensor'> (20L,)\n",
      " |          <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  register_backward_hook(self, hook)\n",
      " |      Registers a backward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time the gradients with respect to module\n",
      " |      inputs are computed. The hook should have the following signature::\n",
      " |      \n",
      " |          hook(module, grad_input, grad_output) -> Tensor or None\n",
      " |      \n",
      " |      The :attr:`grad_input` and :attr:`grad_output` may be tuples if the\n",
      " |      module has multiple inputs or outputs. The hook should not modify its\n",
      " |      arguments, but it can optionally return a new gradient with respect to\n",
      " |      input that will be used in place of :attr:`grad_input` in subsequent\n",
      " |      computations.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |      \n",
      " |      .. warning ::\n",
      " |      \n",
      " |          The current implementation will not have the presented behavior\n",
      " |          for complex :class:`Module` that perform many operations.\n",
      " |          In some failure cases, :attr:`grad_input` and :attr:`grad_output` will only\n",
      " |          contain the gradients for a subset of the inputs and outputs.\n",
      " |          For such :class:`Module`, you should use :func:`torch.Tensor.register_hook`\n",
      " |          directly on a specific input or output to get the required gradients.\n",
      " |  \n",
      " |  register_buffer(self, name, tensor)\n",
      " |      Adds a persistent buffer to the module.\n",
      " |      \n",
      " |      This is typically used to register a buffer that should not to be\n",
      " |      considered a model parameter. For example, BatchNorm's ``running_mean``\n",
      " |      is not a parameter, but is part of the persistent state.\n",
      " |      \n",
      " |      Buffers can be accessed as attributes using given names.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the buffer. The buffer can be accessed\n",
      " |              from this module using the given name\n",
      " |          tensor (Tensor): buffer to be registered.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> self.register_buffer('running_mean', torch.zeros(num_features))\n",
      " |  \n",
      " |  register_forward_hook(self, hook)\n",
      " |      Registers a forward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time after :func:`forward` has computed an output.\n",
      " |      It should have the following signature::\n",
      " |      \n",
      " |          hook(module, input, output) -> None or modified output\n",
      " |      \n",
      " |      The hook can modify the output. It can modify the input inplace but\n",
      " |      it will not have effect on forward since this is called after\n",
      " |      :func:`forward` is called.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_forward_pre_hook(self, hook)\n",
      " |      Registers a forward pre-hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time before :func:`forward` is invoked.\n",
      " |      It should have the following signature::\n",
      " |      \n",
      " |          hook(module, input) -> None or modified input\n",
      " |      \n",
      " |      The hook can modify the input. User can either return a tuple or a\n",
      " |      single modified value in the hook. We will wrap the value into a tuple\n",
      " |      if a single value is returned(unless that value is already a tuple).\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_parameter(self, name, param)\n",
      " |      Adds a parameter to the module.\n",
      " |      \n",
      " |      The parameter can be accessed as an attribute using given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the parameter. The parameter can be accessed\n",
      " |              from this module using the given name\n",
      " |          param (Parameter): parameter to be added to the module.\n",
      " |  \n",
      " |  requires_grad_(self, requires_grad=True)\n",
      " |      Change if autograd should record operations on parameters in this\n",
      " |      module.\n",
      " |      \n",
      " |      This method sets the parameters' :attr:`requires_grad` attributes\n",
      " |      in-place.\n",
      " |      \n",
      " |      This method is helpful for freezing part of the module for finetuning\n",
      " |      or training parts of a model individually (e.g., GAN training).\n",
      " |      \n",
      " |      Args:\n",
      " |          requires_grad (bool): whether autograd should record operations on\n",
      " |                                parameters in this module. Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  share_memory(self)\n",
      " |  \n",
      " |  state_dict(self, destination=None, prefix='', keep_vars=False)\n",
      " |      Returns a dictionary containing a whole state of the module.\n",
      " |      \n",
      " |      Both parameters and persistent buffers (e.g. running averages) are\n",
      " |      included. Keys are corresponding parameter and buffer names.\n",
      " |      \n",
      " |      Returns:\n",
      " |          dict:\n",
      " |              a dictionary containing a whole state of the module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> module.state_dict().keys()\n",
      " |          ['bias', 'weight']\n",
      " |  \n",
      " |  to(self, *args, **kwargs)\n",
      " |      Moves and/or casts the parameters and buffers.\n",
      " |      \n",
      " |      This can be called as\n",
      " |      \n",
      " |      .. function:: to(device=None, dtype=None, non_blocking=False)\n",
      " |      \n",
      " |      .. function:: to(dtype, non_blocking=False)\n",
      " |      \n",
      " |      .. function:: to(tensor, non_blocking=False)\n",
      " |      \n",
      " |      Its signature is similar to :meth:`torch.Tensor.to`, but only accepts\n",
      " |      floating point desired :attr:`dtype` s. In addition, this method will\n",
      " |      only cast the floating point parameters and buffers to :attr:`dtype`\n",
      " |      (if given). The integral parameters and buffers will be moved\n",
      " |      :attr:`device`, if that is given, but with dtypes unchanged. When\n",
      " |      :attr:`non_blocking` is set, it tries to convert/move asynchronously\n",
      " |      with respect to the host if possible, e.g., moving CPU Tensors with\n",
      " |      pinned memory to CUDA devices.\n",
      " |      \n",
      " |      See below for examples.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Args:\n",
      " |          device (:class:`torch.device`): the desired device of the parameters\n",
      " |              and buffers in this module\n",
      " |          dtype (:class:`torch.dtype`): the desired floating point type of\n",
      " |              the floating point parameters and buffers in this module\n",
      " |          tensor (torch.Tensor): Tensor whose dtype and device are the desired\n",
      " |              dtype and device for all parameters and buffers in this module\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> linear = nn.Linear(2, 2)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]])\n",
      " |          >>> linear.to(torch.double)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]], dtype=torch.float64)\n",
      " |          >>> gpu1 = torch.device(\"cuda:1\")\n",
      " |          >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')\n",
      " |          >>> cpu = torch.device(\"cpu\")\n",
      " |          >>> linear.to(cpu)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16)\n",
      " |  \n",
      " |  train(self, mode=True)\n",
      " |      Sets the module in training mode.\n",
      " |      \n",
      " |      This has any effect only on certain modules. See documentations of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      Args:\n",
      " |          mode (bool): whether to set training mode (``True``) or evaluation\n",
      " |                       mode (``False``). Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  type(self, dst_type)\n",
      " |      Casts all parameters and buffers to :attr:`dst_type`.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          dst_type (type or string): the desired type\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  zero_grad(self)\n",
      " |      Sets gradients of all model parameters to zero.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables (if defined)\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object (if defined)\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  dump_patches = False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(tta.ClassificationTTAWrapper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 4.00 GiB total capacity; 2.01 GiB already allocated; 644.80 KiB free; 2.15 GiB reserved in total by PyTorch)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-34-cde8b3243204>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     33\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0mloss_avg\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mavg\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0macc1_avg\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mavg\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0macc5_avg\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mavg\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     34\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 35\u001b[1;33m \u001b[0mtest_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtta_model\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32m<ipython-input-34-cde8b3243204>\u001b[0m in \u001b[0;36mtest_model\u001b[1;34m(model)\u001b[0m\n\u001b[0;32m     18\u001b[0m         \u001b[1;31m# if i%2 == 1:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     19\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 20\u001b[1;33m         \u001b[0mlog_ps\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     21\u001b[0m         \u001b[0mloss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mlog_ps\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;31m#         loss.backward()\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    531\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 532\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    533\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\ttach\\wrappers.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, image, *args)\u001b[0m\n\u001b[0;32m     81\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mtransformer\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtransforms\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     82\u001b[0m             \u001b[0maugmented_image\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mtransformer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0maugment_image\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 83\u001b[1;33m             \u001b[0maugmented_output\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0maugmented_image\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     84\u001b[0m             \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moutput_key\u001b[0m \u001b[1;32mis\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[1;32mNone\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     85\u001b[0m                 \u001b[0maugmented_output\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0maugmented_output\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moutput_key\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    531\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 532\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    533\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torchvision\\models\\resnet.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m    214\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    215\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 216\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_impl\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    217\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    218\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torchvision\\models\\resnet.py\u001b[0m in \u001b[0;36m_forward_impl\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m    204\u001b[0m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlayer1\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    205\u001b[0m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlayer2\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 206\u001b[1;33m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlayer3\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    207\u001b[0m         \u001b[0mx\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mlayer4\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mx\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    208\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    531\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 532\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    533\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\container.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m     98\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     99\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 100\u001b[1;33m             \u001b[0minput\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    101\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    102\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    531\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 532\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    533\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torchvision\\models\\resnet.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m    102\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    103\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 104\u001b[1;33m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconv2\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    105\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mbn2\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    106\u001b[0m         \u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mout\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    531\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 532\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    533\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\conv.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    343\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    344\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 345\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconv2d_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    346\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    347\u001b[0m \u001b[1;32mclass\u001b[0m \u001b[0mConv3d\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0m_ConvNd\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mE:\\soft\\Anaconda3\\envs\\test\\lib\\site-packages\\torch\\nn\\modules\\conv.py\u001b[0m in \u001b[0;36mconv2d_forward\u001b[1;34m(self, input, weight)\u001b[0m\n\u001b[0;32m    340\u001b[0m                             _pair(0), self.dilation, self.groups)\n\u001b[0;32m    341\u001b[0m         return F.conv2d(input, weight, self.bias, self.stride,\n\u001b[1;32m--> 342\u001b[1;33m                         self.padding, self.dilation, self.groups)\n\u001b[0m\u001b[0;32m    343\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    344\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 4.00 GiB total capacity; 2.01 GiB already allocated; 644.80 KiB free; 2.15 GiB reserved in total by PyTorch)"
     ]
    }
   ],
   "source": [
    "def test_model(model):\n",
    "    model = model.cuda()\n",
    "    criterion = set_criterion()\n",
    "    optimizer = set_optimizer(model)\n",
    "    loss_avg = AverageMeter()\n",
    "    acc1_avg= AverageMeter()\n",
    "    acc5_avg= AverageMeter()\n",
    "    data_len = len(test_loader)\n",
    "    \n",
    "    for i, data in enumerate(test_loader):\n",
    "        images, labels = data\n",
    "        if True:\n",
    "            images = Variable(images).cuda()\n",
    "            labels = Variable(labels).cuda()\n",
    "        else:\n",
    "            images = Variable(images)\n",
    "            labels = Variable(labels)\n",
    "        # if i%2 == 1:\n",
    "        optimizer.zero_grad()\n",
    "        log_ps = model(images)\n",
    "        loss = criterion(log_ps, labels)\n",
    "#         loss.backward()\n",
    "        # if i%2 == 1:\n",
    "#         optimizer.step()\n",
    "        print(log_ps.shape)\n",
    "        print(\"loss\", loss)\n",
    "        acc1,  acc5 = get_accuracy(log_ps, labels, topk=(1, 5))\n",
    "        acc1_avg.update(acc1.item(), images.size(0))\n",
    "        acc5_avg.update(acc5.item(), images.size(0))\n",
    "        loss_avg.update(loss.item(), images.size(0))\n",
    "    \n",
    "    print(loss_avg.avg, acc1_avg.avg, acc5_avg.avg)\n",
    "    return loss_avg.avg, acc1_avg.avg, acc5_avg.avg\n",
    "    \n",
    "test_model(tta_model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "215px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
